import pandas as pd
import torch
from transformers import AutoModelForCausalLM, AutoProcessor
from PIL import Image
import numpy as np
import os
from collections import defaultdict
import re
import time

# 创建日志目录
log_dir = "logs"
os.makedirs(log_dir, exist_ok=True)

# 创建日志文件名（带时间戳）
timestamp = time.strftime("%Y%m%d_%H%M%S")
log_file = os.path.join(log_dir, f"event_extraction_results_{timestamp}.log")

# 设置日志函数
def log_message(message, also_print=True):
    """将消息写入日志文件并可选择打印到控制台"""
    with open(log_file, "a", encoding="utf-8") as f:
        f.write(f"{message}\n")
    if also_print:
        print(message)

# 记录开始日志
log_message(f"=== 事件抽取实验开始: {timestamp} ===\n")

# 读取数据集
data_path = r'/mnt/6t/lyh/1/M2E2/M2E2/text/mee.tsv'
log_message(f"正在读取数据文件: {data_path}")

try:
    # 读取tsv文件
    data = pd.read_csv(data_path, sep='\t', encoding='utf-8')
    log_message(f"数据集列名: {data.columns.tolist()}")
except Exception as e:
    log_message(f"读取数据文件出错: {e}")
    # 如果出错，尝试不同的读取方式
    try:
        data = pd.read_csv(data_path, sep='\t', encoding='utf-8', header=0)
        log_message(f"使用替代方法读取成功，数据集列名: {data.columns.tolist()}")
    except Exception as e:
        log_message(f"第二次尝试读取失败: {e}")
        exit(1)

# 加载模型和处理器
model_path = '/mnt/6t/lyh/1/Phi-3.5-vision-instruct'
log_message(f"正在加载模型: {model_path}")
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype="auto", device_map="auto", trust_remote_code=True)
processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
log_message("模型加载完成")

def extract_event_type(text):
    """从模型输出中提取事件类型"""
    # 事件类型选项
    event_options = [
        'Justice:Arrest-Jail',
        'Conflict:Attack',
        'Movement:Transport',
        'Contact:Phone-Write',
        'Life:Die',
        'Conflict:Demonstrate',
        'Transaction:Transfer-Money',
        'Contact:Meet',
        'none'
    ]
    
    # 直接匹配完整事件类型
    for event in event_options:
        if event.lower() in text.lower():
            return event
    
    # 使用正则表达式查找可能的事件类型（格式可能是"事件类型: X"或"事件: X"）
    event_match = re.search(r'(event|type|类型|事件)[:：\s]\s*([a-zA-Z:_\-]+)', text, re.IGNORECASE)
    if event_match:
        potential_event = event_match.group(2).strip()
        # 查找最相似的事件类型
        for event in event_options:
            if event.lower().replace('-', '').replace(':', '') in potential_event.lower().replace('-', '').replace(':', ''):
                return event
    
    # 如果没有匹配到任何事件类型，返回none
    return "none"

def load_image(image_path):
    """加载图像"""
    try:
        if not os.path.exists(image_path):
            log_message(f"图像文件不存在: {image_path}")
            return None
        image = Image.open(image_path).convert("RGB")
        return image
    except Exception as e:
        log_message(f"加载图像出错 {image_path}: {e}")
        return None

def event_extraction_predict(text, image_path):
    """使用Phi-3.5-vision-instruct模型进行事件抽取"""
    # 加载图像
    image = load_image(image_path)
    if image is None:
        return "none", "图像加载失败"
    
    # 准备提示词
    user_content = f"请简单阅读以下文本和观察图像内容，尝试判断是否有事件发生。如果无法确定事件类型或信息不足，请直接输出none。请不要过度推测。\n{text}\n"
    # 准备消息
    messages = [
        {
            "role": "system",
            "content": "根据文本数据和图像事件抽取，可以参考以下事件：'Justice:Arrest-Jail','Conflict:Attack','Movement:Transport','Contact:Phone-Write','Life:Die','Conflict:Demonstrate','Transaction:Transfer-Money','Contact:Meet','none'。尝试判断是否有事件发生。如果无法确定事件类型或信息不足，请直接输出none。请不要过度推测。"
        },
        {
            "role": "user",
            "content": user_content,
        },
        {
            "role": "user",
            "content": f"<|image_1|>\n"
        }
    ]
    
    # 使用处理器生成输入
    prompt = processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = processor(prompt, images=[image], return_tensors="pt").to(model.device)
    
    # 清理GPU内存
    torch.cuda.empty_cache()
    
    # 模型生成输出
    generated_ids = model.generate(**inputs, max_new_tokens=128)
    generated_ids_trimmed = [
        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    
    output_text = processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )
    log_message(f"模型最原始的输出: {output_text}")
    # 处理输出，提取事件类型
    output_text = output_text[0]
    extracted_event = extract_event_type(output_text)
    return extracted_event, output_text

# 初始化计数器
TP = 0
FP = 0
FN = 0
# 每个事件类型的统计
event_stats = defaultdict(lambda: {'TP': 0, 'FP': 0, 'FN': 0, 'support': 0})

# 准备处理数据集
log_message("准备处理数据集...")
examples = []

# 处理数据
for i, row in data.iterrows():
    if i == 0:  # 跳过表头
        continue
    
    try:
        # 提取图像ID、文本和事件类型
        image_id = str(row['#1 ImageID']).strip()
        text = str(row['#2 String']).strip()
        event_type = str(row['#3 String']).strip()
        
        # 验证事件类型
        valid_event_types = [
            'Justice:Arrest-Jail',
            'Conflict:Attack',
            'Movement:Transport',
            'Contact:Phone-Write',
            'Life:Die',
            'Conflict:Demonstrate',
            'Transaction:Transfer-Money',
            'Contact:Meet',
            'none'
        ]
        
        if event_type not in valid_event_types:
            log_message(f"行 {i} 的事件类型 '{event_type}' 无效，设置为'none'")
            event_type = "none"
            
        examples.append({
            'id': i,
            'image_id': image_id,
            'text': text,
            'event_type': event_type
        })
    except Exception as e:
        log_message(f"处理行 {i} 时出错: {e}")

log_message(f"成功准备 {len(examples)} 个样本")
if len(examples) > 0:
    log_message(f"第一个样本: 图像ID={examples[0]['image_id']}, 事件类型={examples[0]['event_type']}")

# 处理样本进行事件抽取
log_message("\n开始处理样本...")
for example in examples:
    log_message(f"\n--- Example {example['id']} ---")
    
    # 构建图像路径
    image_path = f"/mnt/6t/lyh/1/M2E2/M2E2/image/image/image1/{example['image_id']}"
    
    # 记录样本文本
    log_message(f"样本文本: {example['text']}")
    log_message(f"图像路径: {image_path}")
    
    # 进行事件抽取
    predicted_event, raw_output = event_extraction_predict(example['text'], image_path)
    
    true_event = example['event_type']
    log_message(f"真实事件类型: {true_event}")
    log_message(f"预测事件类型: {predicted_event}")
    
    # 记录当前事件类型的支持数
    event_stats[true_event]['support'] += 1
    
    # 评估结果
    if predicted_event == true_event:
        TP += 1
        event_stats[true_event]['TP'] += 1
        result = "TP"
    elif predicted_event != "none" and predicted_event != true_event:
        FP += 1
        event_stats[predicted_event]['FP'] += 1
        result = "FP"
    elif predicted_event == "none" and true_event != "none":
        FN += 1
        event_stats[true_event]['FN'] += 1
        result = "FN"
    else:
        result = "TN"  # True Negative
    
    log_message(f"判断结果: {result}")
    log_message(f"当前统计 - TP: {TP}, FP: {FP}, FN: {FN}")

# 计算总体精确度、召回率和F1值
precision = TP / (TP + FP) if TP + FP > 0 else 0
recall = TP / (TP + FN) if TP + FN > 0 else 0
f1_score = (2 * precision * recall) / (precision + recall) if precision + recall > 0 else 0

# 计算每个事件类型的指标
event_metrics = {}
for event_type, stats in event_stats.items():
    if event_type == "none":
        continue  # 跳过none类型的统计
        
    event_tp = stats['TP']
    event_fp = stats['FP']
    event_fn = stats['FN']
    
    event_precision = event_tp / (event_tp + event_fp) if event_tp + event_fp > 0 else 0
    event_recall = event_tp / (event_tp + event_fn) if event_tp + event_fn > 0 else 0
    event_f1 = (2 * event_precision * event_recall) / (event_precision + event_recall) if event_precision + event_recall > 0 else 0
    
    event_metrics[event_type] = {
        'precision': event_precision,
        'recall': event_recall,
        'f1': event_f1,
        'support': stats['support']
    }

# 输出结果
result_divider = "-" * 80
log_message("\n" + result_divider)
log_message("{:<30} {:<12} {:<12} {:<12} {:<12}".format("事件类型", "精确度", "召回率", "F1值", "支持数"))
log_message(result_divider)

for event_type, metrics in event_metrics.items():
    log_message("{:<30} {:<12.4f} {:<12.4f} {:<12.4f} {:<12}".format(
        event_type,
        metrics['precision'],
        metrics['recall'],
        metrics['f1'],
        metrics['support']
    ))

log_message(result_divider)

# 输出总体结果
log_message(f"\n--- 结果统计 ---")
log_message(f"TP (True Positives): {TP}")
log_message(f"FP (False Positives): {FP}")
log_message(f"FN (False Negatives): {FN}")
log_message(f"精确度 (Precision): {precision:.4f}")
log_message(f"召回率 (Recall): {recall:.4f}")
log_message(f"F1 值 (F1-Score): {f1_score:.4f}")

log_message(f"\n=== 实验结束，日志保存在 {log_file} ===")

